#!/usr/bin/env bash
# Copyright (c) 2022-2024 maminjie <canpool@163.com>
# SPDX-License-Identifier: MulanPSL-2.0

CUR_DIR=$(pwd)
WORK_DIR="$CUR_DIR/build_rootfs"
ROOTFS_DIR="$WORK_DIR/rootfs"
DATE=$(date +%Y%m%d%H%M%S)
OUT_DIR="$CUR_DIR/output"
YUM_DIR="$WORK_DIR/yum"
YUM_CONFIG_FILE="$YUM_DIR/yum.conf"

PRODUCT="openLinux"
VERSION="1.0"
REPOS=""
ROOTFS_NAME=""
PKG_FILES=""

usage() {
printf "build rootfs image

usage:
    bash build_rootfs.sh [-h] [-p PRODUCT] [-v VERSION] [-r RELEASE] [-s REPOSITORY] [-f FILE]

optional arguments:
    -p PRODUCT     product name, default is %s
    -v VERSION     version identifier, default is %s
    -r RELEASE     release information, default is null
    -s REPOSITORY  source dnf repository address link (can be specified multiple times)
    -f FILE        package list file, contains extra packages (can be specified multiple times)
    -h             show the help message and exit
\n" "$PRODUCT" "$VERSION"
}

log_error() {
    echo "ERROR: $1"
}

parse_cmdline() {
    if [ $# -eq 0 ]; then
        usage; exit 0
    fi

    ARCH="$(uname -m)"
    if [[ "${ARCH}" != "aarch64" && "${ARCH}" != "x86_64" ]]; then
        log_error "unsupported architectures: ${ARCH}"; return 1
    fi
    local release=""
    # parse cmd line arguments
    while getopts ":p:v:r:s:f:h" opt; do
        case "$opt" in
            p)
                PRODUCT="$OPTARG"
            ;;
            v)
                VERSION="$OPTARG"
            ;;
            r)
                release="$OPTARG"
            ;;
            s)
                if [ -z "$REPOS" ]; then
                    REPOS="$OPTARG"
                else
                    REPOS="$REPOS $OPTARG"
                fi
            ;;
            f)
                if [ -f "$OPTARG" ]; then
                    if [ -z "$PKG_FILES" ]; then
                        PKG_FILES="$OPTARG"
                    else
                        PKG_FILES="$PKG_FILES $OPTARG"
                    fi
                else
                    log_error "$OPTARG is nof file"; return 1
                fi
            ;;
            h)
                usage; exit 0
            ;;
            ?)
                log_error "please check the params."; usage; return 1
            ;;
        esac
    done
    if [ -z "$REPOS" ]; then
        log_error "invalid REPOSITORY arguments"; return 1
    fi
    if [ -z "$release" ]; then
        ROOTFS_NAME="${PRODUCT}-${VERSION}-${ARCH}-${DATE}.rootfs.ext4"
    else
        ROOTFS_NAME="${PRODUCT}-${VERSION}-${release}-${ARCH}-${DATE}.rootfs.ext4"
    fi
    return 0
}

check_user() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root!"
        exit 1
    fi
}

check_requires() {
    local ret=0

    command -v rsync
    if [ $? -ne 0 ]; then
        log_error "command rsync not found"
        ret=1
    fi

    return $ret
}

clean() {
    cd "$CUR_DIR" || exit
    [[ -d "${WORK_DIR}" ]] && rm -rf "${WORK_DIR}"
}

# create_yum_repo YUM_DIR YUM_CONFIG_FILE REPOS
create_yum_repo() {
    local yum_dir="$1"
    local yum_config_file="$2"
    local repos="$3"
    local yum_repo_dir="${yum_dir}/yum.repos.d"
    local yum_repo_file="${yum_repo_dir}/os.repo"

    mkdir -p "$yum_repo_dir"

    cat > "${yum_config_file}" <<EOF
[main]
gpgcheck=1
installonly_limit=3
clean_requirements_on_remove=True
best=True
skip_if_unavailable=False
reposdir=${yum_repo_dir}
EOF

    local i=1
    for repo in $repos; do
        cat >> "${yum_repo_file}" <<EOF
[os_repo_$i]
name=$repo
baseurl=$repo
enabled=1
gpgcheck=0

EOF
        ((i++))
    done
}

do_main() {
    clean

    mkdir -p "$YUM_DIR"
    create_yum_repo "$YUM_DIR" "$YUM_CONFIG_FILE" "$REPOS"

    mkdir -p $ROOTFS_DIR
    mkdir -p $ROOTFS_DIR/var/lib/rpm
    rpm --root $ROOTFS_DIR --initdb
    # install
    yum install -c "$YUM_CONFIG_FILE" --installroot="$ROOTFS_DIR" --setopt=install_weak_deps=False -y \
        system-release rootfiles vim-minimal passwd hostname \
        NetworkManager openssh-server openssh-clients net-tools iproute iputils ntp dhcp
    if [ $? -ne 0 ]; then
        log_error "yum install basic packages failed"; return 1
    fi
    for f in $PKG_FILES; do
        yum install -c "$YUM_CONFIG_FILE" --installroot="$ROOTFS_DIR" --setopt=install_weak_deps=False -y $(tr '\n' ' ' < "$f")
        if [ $? -ne 0 ]; then
            log_error "yum install customize packages from $f failed"; return 1
        fi
    done
    # TODO: fine cut
    yum clean all -c "$YUM_CONFIG_FILE" --installroot="$ROOTFS_DIR"
    # rm -rf $ROOTFS_DIR/usr/share/{man,doc,locale,icons,info,mime}/*

    # config
    cp /etc/hosts $ROOTFS_DIR/etc/hosts
    cp -L /etc/resolv.conf $ROOTFS_DIR/etc/resolv.conf

    echo "nameserver 8.8.8.8
nameserver 114.114.114.114" > $ROOTFS_DIR/etc/resolv.conf

    if [ ! -d "${ROOTFS_DIR}/etc/sysconfig/network-scripts" ]; then
        mkdir ${ROOTFS_DIR}/etc/sysconfig/network-scripts
    fi

    echo "TYPE=Ethernet
PROXY_METHOD=none
BROWSER_ONLY=no
BOOTPROTO=dhcp
DEFROUTE=yes
IPV4_FAILURE_FATAL=no
IPV6INIT=yes
IPV6_AUTOCONF=yes
IPV6_DEFROUTE=yes
IPV6_FAILURE_FATAL=no
IPV6_ADDR_GEN_MODE=stable-privacy
NAME=eth0
UUID=851a6f36-e65c-3a43-8f4a-78fd0fc09dc9
ONBOOT=yes
AUTOCONNECT_PRIORITY=-999
DEVICE=eth0" > ${ROOTFS_DIR}/etc/sysconfig/network-scripts/ifup-eth0

    # expand the root partition
    cat << EOF > ${ROOTFS_DIR}/etc/rc.d/init.d/expand_rootfs.sh
#!/bin/bash
# chkconfig: - 99 10
ROOT_PART="\$(findmnt / -o source -n)"
resize2fs \$ROOT_PART
if [ -f /etc/rc.d/init.d/expand_rootfs.sh ]; then
    chkconfig --del expand_rootfs.sh
    rm -f /etc/rc.d/init.d/expand_rootfs.sh
fi
EOF
    chmod +x ${ROOTFS_DIR}/etc/rc.d/init.d/expand_rootfs.sh

    cat << EOF | chroot ${ROOTFS_DIR}  /bin/bash
echo "${PRODUCT}#123" | passwd --stdin root
echo $PRODUCT > /etc/hostname
ln -s /usr/share/zoneinfo/Asia/Shanghai /etc/localtime
chkconfig --add expand_rootfs.sh
chkconfig expand_rootfs.sh on
EOF

    local timesyncd_conf="${ROOTFS_DIR}/etc/systemd/timesyncd.conf"
    if [ -f "$timesyncd_conf" ]; then
        sed -i 's/#NTP=/NTP=0.cn.pool.ntp.org/g' $timesyncd_conf
        sed -i 's/#FallbackNTP=/FallbackNTP=1.asia.pool.ntp.org 2.asia.pool.ntp.org/g' $timesyncd_conf
    fi

    if [ ! -d "$OUT_DIR" ]; then
        mkdir -p $OUT_DIR
    fi

    # make image
    cd $WORK_DIR || exit

    local size=$(du -sh --block-size=1MiB ${ROOTFS_DIR} | cut -f 1 | xargs)
    size=$(($size+500))
    rootfs_img=${OUT_DIR}/${ROOTFS_NAME}
    dd if=/dev/zero of=${rootfs_img} bs=1MiB count=$size status=progress && sync
    mkfs.ext4 -L rootfs ${rootfs_img}

    local rootfs_tmpdir=${WORK_DIR}/rootfs_tmp
    mkdir ${rootfs_tmpdir}
    mount ${rootfs_img} ${rootfs_tmpdir}

    rsync -avHAXq ${ROOTFS_DIR}/* ${rootfs_tmpdir}
    sync
    sleep 2

    if [ -d ${rootfs_tmpdir} ]; then
        if grep -q "${rootfs_tmpdir}" /proc/mounts ; then
            umount ${rootfs_tmpdir}
        fi
    fi

    if [ ! -f $rootfs_img ]; then
        return 1
    fi

    clean

    return 0
}

main() {
    check_user

    parse_cmdline "$@"
    if [ $? -ne 0 ]; then
        log_error "parse params failed, try -h for help"
        return 1
    fi

    check_requires
    if [ $? -ne 0 ]; then
        log_error "check requires failed"
        return 1
    fi

    do_main
    if [ $? -ne 0 ]; then
        return 1
    fi

    return 0
}

main "$@"
